package com.zqq.demo.camera1;

import android.media.MediaCodec;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import java.io.IOException;
import java.nio.ByteBuffer;

public class H264ToMp4Converter {
    private MediaMuxer mMediaMuxer;
    private int mVideoTrackIndex;
    private boolean mMuxerStarted;
    
    public void init(String outputPath, MediaFormat format,int width, int height) throws IOException {
        mMediaMuxer = new MediaMuxer(outputPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
        mVideoTrackIndex = mMediaMuxer.addTrack(format);
        mMuxerStarted = false;
    }
    private long mLastPresentationTimeUs = 0;
    public void writeSampleData(byte[] h264Data, long presentationTimeUs) {
        if (!mMuxerStarted) {
            mMediaMuxer.start();
            mMuxerStarted = true;
        }
        
        ByteBuffer buffer = ByteBuffer.wrap(h264Data);
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
//        bufferInfo.set(0, h264Data.length, presentationTimeUs,
//            (h264Data[4] & 0x1F) == 0x05 ? MediaCodec.BUFFER_FLAG_KEY_FRAME : 0);
        // 计算时间戳（30fps）
        mLastPresentationTimeUs += 1000000 / 30;
        bufferInfo.presentationTimeUs = mLastPresentationTimeUs;
        bufferInfo.size = h264Data.length;
        bufferInfo.flags = (h264Data[4] & 0x1F) == 0x05 ?
                MediaCodec.BUFFER_FLAG_KEY_FRAME : 0;

        mMediaMuxer.writeSampleData(mVideoTrackIndex, buffer, bufferInfo);
    }

    private long mFrameIntervalUs = 1000000 / 30;
    private long mLastPtsUs = 0;
    public void writeFrame(byte[] frameData, boolean isKeyFrame) {
        if (!mMuxerStarted && isKeyFrame) {
            mMediaMuxer.start();
            mMuxerStarted = true;
        }

        if (mMuxerStarted) {
            ByteBuffer buffer = ByteBuffer.wrap(frameData);
            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            bufferInfo.offset = 0;
            bufferInfo.size = frameData.length;
            bufferInfo.presentationTimeUs = mLastPtsUs;
            bufferInfo.flags = isKeyFrame ?
                    MediaCodec.BUFFER_FLAG_KEY_FRAME : 0;

            mMediaMuxer.writeSampleData(mVideoTrackIndex, buffer, bufferInfo);
            mLastPtsUs += mFrameIntervalUs;
        }
    }

    public void release() {
        if (mMediaMuxer != null) {
            try {
                if (mMuxerStarted) {
                    mMediaMuxer.stop();
                }
                mMediaMuxer.release();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
